#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import socketserver
import threading
import queue
from typing import Optional, Callable, List, Tuple
import socket


def get_host_ip():
    """返回局域网可用的 IPv4 地址"""
    try:
        # s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        # s.connect(("8.8.8.8", 80))
        # ip = s.getsockname()[0]
        # s.close()

        ip = socket.gethostbyname(socket.gethostname())

    except Exception:
        ip = "127.0.0.1"

    print(ip)
    return ip

HOST = get_host_ip()

# HOST = "0.0.0.0"
PORT = 6002

# ======== 客户端记录与消息队列 ========
CLIENTS = set()  # 保存 (ip, port)
CLIENTS_LOCK = threading.Lock()

INBOX: "queue.Queue[Tuple[str, bytes]]" = queue.Queue()
_ON_MESSAGE: Optional[Callable[[str, bytes], None]] = None
_ON_MESSAGE_LOCK = threading.Lock()

def _emit_message(client_id: str, data: bytes):
    INBOX.put((client_id, data))
    with _ON_MESSAGE_LOCK:
        cb = _ON_MESSAGE
    if cb:
        try:
            cb(client_id, data)
        except Exception as e:
            print(f"[WARN] on_message callback error: {e}")

# ======== 外部可调用接口 ========
def set_on_message(cb: Optional[Callable[[str, bytes], None]]):
    """设置或取消收到消息的回调函数：callback(client_id, data_bytes)"""
    global _ON_MESSAGE
    with _ON_MESSAGE_LOCK:
        _ON_MESSAGE = cb

def recv_next(timeout: Optional[float] = None) -> Tuple[str, bytes]:
    """阻塞等待一条消息"""
    return INBOX.get(timeout=timeout)

def recv_all_nowait() -> List[Tuple[str, bytes]]:
    """一次性取出所有当前待处理消息"""
    out = []
    while True:
        try:
            out.append(INBOX.get_nowait())
        except queue.Empty:
            break
    return out

# ======== UDP Handler ========
class MyUDPHandler(socketserver.BaseRequestHandler):
    def handle(self):
        data, sock = self.request  # data: bytes, sock: socket.socket
        client_ip, client_port = self.client_address
        client_id = f"{client_ip}:{client_port}"

        # 记录客户端
        with CLIENTS_LOCK:
            CLIENTS.add((client_ip, client_port))

        # print(f"[RX {client_id}] {data.decode('utf-8', 'ignore').rstrip()}")

        # 加入队列 & 回调
        _emit_message(client_id, data)

        # 回一条确认信息（可选）
        # ack = f"ACK from server, got {len(data)} bytes".encode("utf-8")
        # sock.sendto(ack, self.client_address)

# ======== 对外发送接口 ========
def send_text_to(client_id: str, text: str):
    """向指定客户端发送文本"""
    host, port_str = client_id.split(":")
    port = int(port_str)
    with socketserver.socket.socket(socketserver.socket.AF_INET, socketserver.socket.SOCK_DGRAM) as s:
        s.sendto(text.encode("utf-8"), (host, port))

def send_text_to_all(text: str):
    """向所有已知客户端群发"""
    failed = []
    with CLIENTS_LOCK:
        targets = list(CLIENTS)
        print(targets)
    for ip, port in targets:
        try:
            with socketserver.socket.socket(socketserver.socket.AF_INET, socketserver.socket.SOCK_DGRAM) as s:
                s.sendto(text.encode("utf-8"), (ip, port))
        except Exception as e:
            failed.append((f"{ip}:{port}", str(e)))
    return failed

# ======== 可嵌入的启动/停止接口 ========
_SERVER = None
_SERVER_THREAD = None
_SERVER_LOCK = threading.Lock()

def start_server(host=HOST, port=PORT):
    """启动 UDP 服务器"""
    global _SERVER, _SERVER_THREAD
    with _SERVER_LOCK:
        if _SERVER is not None:
            return _SERVER.server_address
        _SERVER = socketserver.ThreadingUDPServer((host, port), MyUDPHandler)
        _SERVER.allow_reuse_address = True

        def _serve():
            try:
                _SERVER.serve_forever()
            except Exception as e:
                print(f"[SERV] serve_forever exited: {e}")

        _SERVER_THREAD = threading.Thread(target=_serve, daemon=True)
        _SERVER_THREAD.start()
        print(f"[SERV] UDP listening on {host}:{port}")
        return _SERVER.server_address

def stop_server():
    """停止 UDP 服务器"""
    global _SERVER, _SERVER_THREAD
    with _SERVER_LOCK:
        if _SERVER is None:
            return
        try:
            _SERVER.shutdown()
            _SERVER.server_close()
        finally:
            _SERVER = None
        if _SERVER_THREAD:
            _SERVER_THREAD.join(timeout=1.0)
            _SERVER_THREAD = None
        print("[SERV] UDP server stopped")

# ======== 测试入口 ========
if __name__ == "__main__":
    start_server()
    print("[INFO] Server running")
    try:
        while True:
            cid, data = recv_next()
            # print(f"[INBOX] {cid}: {data!r}")
    except KeyboardInterrupt:
        stop_server()
